package com.example.terminalservice.utils;

import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * @author 28142
 * @description ssh连接池工具类
 * @date 2023/8/7 22:06
 */
public class SSHUtils {
    private static JSch jsch;
    private static ConcurrentMap<String, GenericObjectPool<Session>> sessionPools;

    static {
        jsch = new JSch();
        sessionPools = new ConcurrentHashMap<>();
    }

    public static Session getSession(String host, int port, String username, String password) throws Exception {
        System.out.println("*********调用getSession*********");
        String sessionKey = getSessionKey(host, port, username);

        // 先尝试从缓存中获取连接池
        GenericObjectPool<Session> sessionPool = sessionPools.get(sessionKey);
        if (sessionPool == null) {
            // 如果连接池不存在，则创建连接池并放入缓存
            sessionPool = createSessionPool(host, port, username, password);
            sessionPools.put(sessionKey, sessionPool);
        }
        return sessionPool.borrowObject();
    }

    public static void returnSession(Session session) {
        if (session != null) {
            String sessionKey = getSessionKey(session);
            GenericObjectPool<Session> sessionPool = sessionPools.get(sessionKey);
            if (sessionPool != null) {
                sessionPool.returnObject(session);
            } else {
                // Handle the case when the session pool is not found (e.g., closed or removed)
                session.disconnect();
            }
        }
    }

    private static String getSessionKey(Session session) {
        return getSessionKey(session.getHost(), session.getPort(), session.getUserName());
    }

    private static String getSessionKey(String host, int port, String username) {
        return host + ":" + port + ":" + username;
    }

    private static GenericObjectPool<Session> createSessionPool(String host, int port, String username, String password) {
        System.out.println("================调用createSessionPool方法================");
        GenericObjectPoolConfig<Session> poolConfig = new GenericObjectPoolConfig<>();
        poolConfig.setMaxTotal(10);
        poolConfig.setMaxIdle(5);
        poolConfig.setMinIdle(1);

        return new GenericObjectPool<>(new SessionFactory(host, port, username, password), poolConfig);
    }

    private static class SessionFactory extends BasePooledObjectFactory<Session> {
        private String host;
        private int port;
        private String username;
        private String password;

        public SessionFactory(String host, int port, String username, String password) {
            this.host = host;
            this.port = port;
            this.username = username;
            this.password = password;
        }

        @Override
        public Session create() throws Exception {
            System.out.println("******************调用create方法******************");
            Session session = jsch.getSession(username, host, port);
            session.setPassword(password);

            // 设置不进行主机密钥检查（仅在测试环境中使用，生产环境请慎用）
            Properties config = new Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);

            // 设置连接超时，单位是毫秒，比如设置为10秒
            int timeout = 10000;

            session.connect(timeout);
            return session;
        }

        @Override
        public PooledObject<Session> wrap(Session obj) {
            return new DefaultPooledObject<>(obj);
        }

        @Override
        public void destroyObject(PooledObject<Session> p) throws Exception {
            Session session = p.getObject();
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
        }
    }
}
